from torch import nn

from utils.model_utils import get_polygon_encoder


class POLYGON_EMBEDDING_MODEL(nn.Module):
    def __init__(self, args, spa_enc, geom_type, extent):
        super(POLYGON_EMBEDDING_MODEL, self).__init__()
        self.pgon_enc = get_polygon_encoder(args,
                                            pgon_enc_type=args.pgon_enc,
                                            spa_enc=spa_enc,
                                            pgon_embed_dim=args.pgon_embed_dim,
                                            extent=extent[geom_type],
                                            padding_mode=args.padding_mode,
                                            dropout=args.dropout,
                                            device=args.device)

    def compute_polygon_embed(self, polygons):
        polygons_seq = polygons[:, :-1, :]
        pgon_embeds = self.pgon_enc(polygons_seq)
        return pgon_embeds

    def forward(self, polygons):
        pgon_embeds = self.compute_polygon_embed(polygons)
        return pgon_embeds
